package lion.dev.encrypt.cert;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringReader;
import java.math.BigInteger;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Security;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.output.StringBuilderWriter;
import org.apache.commons.lang.StringUtils;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.KeyUsage;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x509.X509Extension;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.pkcs.PKCS10CertificationRequest;
import org.bouncycastle.pkcs.PKCS10CertificationRequestBuilder;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemReader;
import org.bouncycastle.util.io.pem.PemWriter;

import sun.security.x509.AlgorithmId;
import sun.security.x509.X509CertImpl;
import sun.security.x509.X509CertInfo;

/**
 * @author lion
 * @email hl_0074@sina.com
 * @date 2016年3月2日
 */
@SuppressWarnings({ "restriction", "deprecation" })
public class JKStore {

	static {
		Security.addProvider(new BouncyCastleProvider());
	}

	public static final String PEM_TYPE_CERTIFICATE = "CERTIFICATE";
	public static final String PEM_TYPE_CERTIFICATE_REQUEST = "CERTIFICATE REQUEST";

	public static final String KEYSTORE_TYPE_JKS = "JKS";
	public static final String KEYSTORE_TYPE_PKCS12 = "PKCS12";

	public static final String CERT_TYPE = "x.509";
	public static final String DEFAULT_SIGN_ALG = "SHA256withRSA";

	private String passwd;
	private String path;
	private String type = KeyStore.getDefaultType();

	private KeyStore keyStore;

	public JKStore(String path, String passwd, String type) {
		this.path = path;
		this.passwd = passwd;
		this.type = type;
		try (FileInputStream fis = new FileInputStream(path)) {
			this.keyStore = KeyStore.getInstance(type);
			this.keyStore.load(fis, this.passwd == null ? null : this.passwd.toCharArray());
		} catch (IOException e) { // 如果路径中没有密钥库，则创建一个默认的
			e.printStackTrace();
			try {
				this.keyStore = KeyStore.getInstance(type);
				this.keyStore.load(null, null);
				this.update();
			} catch (Exception e1) {
			}
		} catch (Exception e) {
		}
	}

	public JKStore(String path, String passwd) {
		this(path, passwd, KeyStore.getDefaultType());
	}

	public boolean genKeyPair(String alias, String passwd, CertInfo info) {

		if (passwd == null) {
			passwd = "";
		}

		try {
			KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
			kpg.initialize(2048);

			KeyPair keyPair = kpg.generateKeyPair();
			X509Certificate certificate = signCert(keyPair.getPrivate(),
					SubjectPublicKeyInfo.getInstance(keyPair.getPublic()), info);

			this.keyStore.setKeyEntry(alias, keyPair.getPrivate(), passwd.toCharArray(),
					new Certificate[] { certificate });

			this.update();
			return true;
		} catch (Exception e) {
			e.printStackTrace();
		}
		return false;
	}

	/**
	 * import a cert
	 *
	 * @param alias
	 *            cert entry alias
	 * @param cert
	 *            cert binary content
	 */
	public boolean importCert(String alias, byte[] cert) {

		return importCert(alias, new ByteArrayInputStream(cert));
	}

	/**
	 * import a cert
	 *
	 * @param alias
	 *            cert entry alias
	 * @param cert
	 *            cert inputstream
	 */
	public boolean importCert(String alias, InputStream stream) {

		try {
			X509Certificate certificate = loadCertificate(stream);
			KeyStore.TrustedCertificateEntry trustedEntry = new KeyStore.TrustedCertificateEntry(certificate);
			keyStore.setEntry(alias, trustedEntry, null);
			return true;
		} catch (KeyStoreException | CertificateException e) {
			e.printStackTrace();
		}
		return false;
	}

	/**
	 * import a cert
	 *
	 * @param alias
	 *            cert entry alias
	 * @param cert
	 *            cert file path
	 */
	public boolean importCert(String alias, String cert) {

		try {
			return importCert(alias, new FileInputStream(cert));
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		}
		return false;
	}

	public byte[] exportKey(String alias, String password) throws Exception {

		Key key = this.keyStore.getKey(alias, password.toCharArray());
		return key.getEncoded();
	}

	/**
	 * export as .p12 file
	 */
	public byte[] exportP12(String alias, String passwd) throws Exception {

		Certificate cert = this.keyStore.getCertificate(alias);

		X509CertImpl cimp = new X509CertImpl(cert.getEncoded());
		X509CertInfo cinfo = (X509CertInfo) cimp.get("x509.info");
		cinfo.set("algorithmID.algorithm", new AlgorithmId(AlgorithmId.sha1WithRSAEncryption_oid));
		X509CertImpl newcert = new X509CertImpl(cinfo);

		char[] pass = passwd == null ? null : passwd.toCharArray();

		PrivateKey prikey = (PrivateKey) this.keyStore.getKey(alias, pass);
		newcert.sign(prikey, DEFAULT_SIGN_ALG);

		Certificate[] cchain = { newcert };
		KeyStore store = KeyStore.getInstance(KEYSTORE_TYPE_PKCS12);
		store.load(null, null);
		store.setKeyEntry(alias, prikey, pass, cchain);

		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		store.store(baos, pass);

		return baos.toByteArray();
	}

	/**
	 * sign a cert
	 *
	 * @param alias
	 *            private key entry alias
	 * @param passwd
	 *            private key entry passwd
	 * @param certPath
	 *            cert file path
	 */
	public byte[] signCert(String alias, String passwd, String cert, CertInfo certInfo) throws Exception {

		char[] pass = passwd == null ? null : passwd.toCharArray();
		PrivateKey privateKey = (PrivateKey) this.keyStore.getKey(alias, pass);
		// cert 2
		X509Certificate certificate2 = loadCertificate(cert);

		return signCert(privateKey, SubjectPublicKeyInfo.getInstance(certificate2.getPublicKey()), certInfo)
				.getEncoded();
	}

	/**
	 * sign a cert
	 *
	 * @param alias
	 *            private key entry alias
	 * @param passwd
	 *            private key entry passwd
	 * @param cert
	 *            cert binary content
	 */
	public byte[] signCert(String alias, String passwd, byte[] cert, CertInfo info) throws Exception {

		char[] pass = passwd == null ? null : passwd.toCharArray();
		PrivateKey privateKey = (PrivateKey) this.keyStore.getKey(alias, pass);
		// cert 2
		X509Certificate clientCert = loadCertificate(cert);
		return signCert(privateKey, SubjectPublicKeyInfo.getInstance(clientCert.getPublicKey()), info).getEncoded();
	}

	/**
	 * sign cert
	 *
	 * @param privateKey
	 *            PrivateKey
	 * @param spki
	 *            SubjectPublicKeyInfo
	 * @param info
	 *            CertInfo
	 */
	public static X509Certificate signCert(PrivateKey privateKey, SubjectPublicKeyInfo spki, CertInfo info)
			throws Exception {

		X509v3CertificateBuilder builder = new X509v3CertificateBuilder(new X500Name(info.getIssuerDN().toString()),
				BigInteger.valueOf(info.getSerialNumber()), info.getNotBefore(), info.getNotAfter(),
				new X500Name(info.getSubjectDN().toString()), spki);

		if (info.getKeyUsage() > 0) {
			builder.addExtension(X509Extension.keyUsage, false, new KeyUsage(info.getKeyUsage()));
		}

		String signAlg = info.getSignAlgName();
		if (StringUtils.isBlank(signAlg)) {
			signAlg = DEFAULT_SIGN_ALG;
		}
		ContentSigner sigGen = new JcaContentSignerBuilder(signAlg).setProvider(BouncyCastleProvider.PROVIDER_NAME)
				.build(privateKey);
		X509CertificateHolder holder = builder.build(sigGen);
		CertificateFactory cf = CertificateFactory.getInstance(CERT_TYPE);

		ByteArrayInputStream is1 = new ByteArrayInputStream(holder.toASN1Structure().getEncoded());
		X509Certificate theCert = (X509Certificate) cf.generateCertificate(is1);

		return theCert;
	}

	/**
	 * verify a certificate
	 *
	 * @param alias
	 *            ca cert alias
	 * @param clientCert
	 *            client cert file path
	 */
	public boolean verify(String alias, String clientCert) {

		try {
			Certificate caCert = this.keyStore.getCertificate(alias);
			Certificate clientCer = loadCertificate(clientCert);
			clientCer.verify(caCert.getPublicKey());
			return true;
		} catch (Exception e) {
			e.printStackTrace();
		}
		return false;
	}

	/**
	 * verify a certificate
	 *
	 * @param alias
	 *            ca cert alias
	 * @param clientCert
	 *            client cert file path
	 */
	public boolean verify(String alias, byte[] clientCert) {

		try {
			Certificate caCert = this.keyStore.getCertificate(alias);
			Certificate clientCer = loadCertificate(clientCert);
			clientCer.verify(caCert.getPublicKey());
			return true;
		} catch (Exception e) {
			e.printStackTrace();
		}
		return false;
	}

	/**
	 * verify a certificate
	 *
	 * @param cacert
	 *            cacert file path
	 * @param clientCert
	 *            clientcert file path
	 *
	 */
	public static boolean verifyCert(String cacert, String clientCert) {

		try {
			Certificate certificate = loadCertificate(cacert);
			Certificate clientCer = loadCertificate(clientCert);
			clientCer.verify(certificate.getPublicKey());
			return true;
		} catch (Exception e) {
		}
		return false;
	}

	/**
	 * export certificate
	 */
	public byte[] exportCert(String alias) throws KeyStoreException, CertificateEncodingException {

		Certificate certificate = this.keyStore.getCertificate(alias);
		return certificate.getEncoded();

	}

	/**
	 * update the keystore
	 */
	public void update() {

		if (StringUtils.isNotBlank(this.path)) {
			try (FileOutputStream fis = new FileOutputStream(this.path)) {
				this.keyStore.store(fis, this.passwd.toCharArray());
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
	}

	/**
	 * create certificate request .csr
	 */
	public byte[] genCSR(String alias, String password) throws Exception {

		X509Certificate certificate = (X509Certificate) this.keyStore.getCertificate(alias);

		PKCS10CertificationRequestBuilder builder = new PKCS10CertificationRequestBuilder(
				new X500Name(certificate.getSubjectDN().getName()),
				SubjectPublicKeyInfo.getInstance(certificate.getPublicKey()));

		JcaContentSignerBuilder jcaContentSignerBuilder = new JcaContentSignerBuilder(certificate.getSigAlgName());
		jcaContentSignerBuilder.setProvider(BouncyCastleProvider.PROVIDER_NAME);
		ContentSigner contentSigner = jcaContentSignerBuilder
				.build((PrivateKey) this.keyStore.getKey(alias, password == null ? null : password.toCharArray()));

		PKCS10CertificationRequest request = builder.build(contentSigner);

		return request.getEncoded();
	}

	/**
	 * sign a certificate request
	 *
	 * @throws Exception
	 */
	public byte[] signCSR(String alias, String password, String csrpath, CertInfo info) throws Exception {

		String pemContent = FileUtils.readFileToString(new File(csrpath));
		byte[] csrContent = readPEMContent(pemContent);

		return signCSR(alias, password, csrContent, info);
	}

	/**
	 * sign a certificate request
	 *
	 * @throws Exception
	 */
	public byte[] signCSR(String alias, String password, byte[] csrContent, CertInfo info) throws Exception {

		PrivateKey privateKey = (PrivateKey) this.keyStore.getKey(alias,
				password == null ? null : password.toCharArray());

		PKCS10CertificationRequest csr = new PKCS10CertificationRequest(csrContent);

		X509Certificate certificate = signCert(privateKey, csr.getSubjectPublicKeyInfo(), info);

		return certificate.getEncoded();
	}

	/**
	 * delete item
	 */
	public boolean delete(String alias) {

		try {
			this.keyStore.deleteEntry(alias);
			return true;
		} catch (KeyStoreException e) {
		}
		return false;
	}

	/**
	 * change entry password
	 */
	public boolean setPasswd(String alias, String passwd, String newPasswd) {

		try {
			Certificate[] chain = this.keyStore.getCertificateChain(alias);
			PrivateKey key = (PrivateKey) this.keyStore.getKey(alias, passwd == null ? null : passwd.toCharArray());

			this.keyStore.setKeyEntry(alias, key, newPasswd == null ? null : newPasswd.toCharArray(), chain);
			return true;
		} catch (UnrecoverableKeyException | KeyStoreException | NoSuchAlgorithmException e) {
			e.printStackTrace();
		}
		return false;
	}

	/**
	 * load loadCertificate from a file path
	 *
	 * @param path
	 *            .cer or .der file path
	 * @throws IOException
	 *
	 */
	public static X509Certificate loadCertificate(String path) throws CertificateException, IOException {

		String text = FileUtils.readFileToString(new File(path));
		if (StringUtils.contains(text, "BEGIN")) { return loadCertificate(readPEMContent(text)); }

		return loadCertificate(new FileInputStream(path));
	}

	/**
	 * load loadCertificate from a byte array
	 *
	 * @param cert
	 *            certificate binary content
	 *
	 */
	public static X509Certificate loadCertificate(byte[] cert) throws CertificateException {

		return loadCertificate(new ByteArrayInputStream(cert));
	}

	/**
	 * load loadCertificate from a stream
	 */
	public static X509Certificate loadCertificate(InputStream stream) throws CertificateException {

		try {
			CertificateFactory certificateFactory = CertificateFactory.getInstance(CERT_TYPE);
			return (X509Certificate) certificateFactory.generateCertificate(stream);
		} finally {
			try {
				stream.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}

	/**
	 * convert to PEM format string
	 *
	 * @param type
	 *            PEM_TYPE_CERTIFICATE or PEM_TYPE_CERTIFICATE_REQUEST
	 * @param content
	 *            the content to convert
	 */
	public static String toPEMString(String type, byte[] content) {

		StringBuilderWriter sbw = new StringBuilderWriter();
		PemWriter writer = new PemWriter(sbw);
		PemObject pemObject = new PemObject(type, content);
		try {
			writer.writeObject(pemObject);
			writer.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
		return sbw.toString();
	}

	/**
	 * read content from a pem format string
	 *
	 * @param pemContent
	 *            pem string content
	 */
	public static byte[] readPEMContent(String pemContent) throws IOException {

		try (PemReader reader = new PemReader(new StringReader(pemContent))) {
			return reader.readPemObject().getContent();
		}
	}

	public String getPasswd() {

		return passwd;
	}

	public void setPasswd(String passwd) {

		this.passwd = passwd;
	}

	public KeyStore getKeyStore() {

		return keyStore;
	}

	public String getType() {

		return type;
	}
}
